import json
import requests
import time
from typing import List, Dict, Any

from volcengine.auth.SignerV4 import SignerV4
from volcengine.base.Request import Request
from volcengine.Credentials import Credentials

account_id = "2104289245"
g_knowledge_base_domain = "api-knowledgebase.mlp.cn-beijing.volces.com"
apikey = "507a0b8d-b93d-48fe-844c-96205a2d8b48"


def prepare_request(method, path, params=None, data=None, doseq=0):
    if params:
        for key in params:
            if (
                    isinstance(params[key], int)
                    or isinstance(params[key], float)
                    or isinstance(params[key], bool)
            ):
                params[key] = str(params[key])
            elif isinstance(params[key], list):
                if not doseq:
                    params[key] = ",".join(params[key])
    r = Request()
    r.set_shema("http")
    r.set_method(method)
    r.set_connection_timeout(10)
    r.set_socket_timeout(10)
    headers = {
        "Accept": "application/json",
        "Content-Type": "application/json;charset=UTF-8",
        "Host": g_knowledge_base_domain,
        'Authorization': f'Bearer {apikey}'
    }
    r.set_headers(headers)
    if params:
        r.set_query(params)
    r.set_host(g_knowledge_base_domain)
    r.set_path(path)
    if data is not None:
        r.set_body(json.dumps(data))
    return r


def knowledge_service_chat(query):
    """
    调用知识库服务查询单个查询
    """
    method = "POST"
    path = "/api/knowledge/service/chat"
    request_params = {
        "service_resource_id": "kb-service-379a654af739335c",
        "messages": [
            {
                "role": "user",
                "content": query
            }
        ],
        "stream": False
    }

    try:
        info_req = prepare_request(method=method, path=path, data=request_params)
        rsp = requests.request(
            method=info_req.method,
            url="http://{}{}".format(g_knowledge_base_domain, info_req.path),
            headers=info_req.headers,
            data=info_req.body
        )
        rsp.encoding = "utf-8"
        response_data = json.loads(rsp.text)
        
        if response_data.get("code") == 0 and "data" in response_data:
            result_list = response_data["data"].get("result_list", [])
            
            # Find the item with recall_position = 1
            for item in result_list:
                if item.get("recall_position") == 1:
                    table_chunk_fields = item.get("table_chunk_fields", [])
                    
                    # Format the result as a structured object
                    if table_chunk_fields:
                        result_obj = {}
                        
                        # 遍历所有字段，找到name
                        for field in table_chunk_fields:
                            field_name = field.get("field_name")
                            field_value = field.get("field_value")
                            if field_name == "name" and field_value:
                                result_obj["name"] = field_value
                                return result_obj
                        
                        # 如果没有找到name，返回空对象
                        return {"name": ""}
            
            # 如果没有找到recall_position = 1的结果，返回空对象
            return {"name": ""}
        else:
            error_msg = response_data.get("message", "Unknown error")
            return {"error": f"Material specification query failed: {error_msg}"}
            
    except json.JSONDecodeError as e:
        error_msg = f"JSON解析错误: {str(e)}"
        return {"error": error_msg}
    except Exception as e:
        error_msg = f"请求失败: {str(e)}"
        return {"error": error_msg}


def batch_knowledge_service_chat(queries: List[str], delay: float = 0.5) -> Dict[str, List[Dict[str, str]]]:
    """
    批量调用知识库服务查询
    
    Args:
        queries: 查询列表
        delay: 请求间隔时间（秒），避免请求过快
        
    Returns:
        Dict格式的查询结果:
        {
            "data": [
                {
                    "查询内容": str,
                    "name": str
                },
                ...
            ]
        }
    """
    results = []
    
    print(f"开始批量查询，共 {len(queries)} 个查询")
    
    for i, query in enumerate(queries, 1):
        print(f"处理第 {i}/{len(queries)} 个查询: {query[:50]}...")
        
        try:
            # 调用单个查询
            result = knowledge_service_chat(query)
            
            # 构建标准格式的结果
            formatted_result = {
                
                query : result.get("name", "")  # 从结果中获取name作为name字段
            }
            
            results.append(formatted_result)
            
            # 显示查询结果
            if "error" in result:
                print(f"  查询失败: {result['error']}")
            else:
                print(f"  查询成功: {formatted_result}")
            
            # 添加延迟避免请求过快
            if i < len(queries):  # 最后一个查询不需要延迟
                time.sleep(delay)
                
        except Exception as e:
            print(f"  查询异常: {str(e)}")
            results.append({
                "查询内容": query,
                "name": f"查询异常: {str(e)}"
            })
    
    # 返回标准格式的结果
    return {"data": results}


if __name__ == "__main__":

    # 测试批量查询
    print("\n=== 批量查询测试 ===")
    test_queries = [
        "CZN 3+",
        "ZnB3+",
        "三价铬黑锌",
        "环保彩锌",
        "锌镍合金本色"
    ]
    
    batch_results = batch_knowledge_service_chat(test_queries, delay=0.5)
    